#!/usr/bin/env python3
"""
Analyze the business-payout.pkl graph data to understand node structure
"""

import pickle
import json
from collections import defaultdict

# Load the graph
pkl_file = "/data_ext/ref_code/LocAgent/index_data_strict/pp_data/graph_index_v2.3/business-payout.pkl"

print("Loading graph from pickle file...")
with open(pkl_file, "rb") as f:
    G = pickle.load(f)

print(f"\n=== Graph Statistics ===")
print(f"Total nodes: {G.number_of_nodes()}")
print(f"Total edges: {G.number_of_edges()}")

# Analyze node types
node_types = defaultdict(int)
nodes_with_code = 0
nodes_without_code = 0
sample_nodes = {}

print(f"\n=== Analyzing Node Structure ===")
for node_id in list(G.nodes())[:100]:  # Sample first 100 nodes
    node_data = G.nodes[node_id]

    # Count node types
    if 'type' in node_data:
        node_type = node_data['type']
        node_types[node_type] += 1

        # Save sample for each type
        if node_type not in sample_nodes:
            sample_nodes[node_type] = {
                'node_id': node_id,
                'data': node_data
            }

    # Check if code_content exists
    if 'code_content' in node_data:
        nodes_with_code += 1
    else:
        nodes_without_code += 1

print(f"\nNodes with code_content: {nodes_with_code}")
print(f"Nodes without code_content: {nodes_without_code}")

# Show node type distribution (from first 100 nodes)
print(f"\n=== Node Type Distribution (sample) ===")
for ntype, count in sorted(node_types.items()):
    print(f"{ntype}: {count}")

# Display sample node structure for each type
print(f"\n=== Sample Node Structure for Each Type ===")
for node_type, sample in sample_nodes.items():
    print(f"\n--- {node_type.upper()} Node Sample ---")
    print(f"Node ID: {sample['node_id']}")

    # Show all attributes except code_content (too long)
    for key, value in sample['data'].items():
        if key == 'code_content':
            if value:
                lines = value.split('\n')
                print(f"  {key}: <{len(lines)} lines, {len(value)} chars>")
                # Show first 3 lines as preview
                preview = '\n'.join(lines[:3])
                if len(preview) > 200:
                    preview = preview[:200] + "..."
                print(f"    Preview: {repr(preview)}")
            else:
                print(f"  {key}: None")
        else:
            if isinstance(value, (str, int, float, bool, type(None))):
                print(f"  {key}: {value}")
            elif isinstance(value, (list, dict)):
                print(f"  {key}: {type(value).__name__} with {len(value)} items")
            else:
                print(f"  {key}: {type(value).__name__}")

# Analyze edges
print(f"\n=== Edge Type Distribution ===")
edge_types = defaultdict(int)
for u, v, data in list(G.edges(data=True))[:100]:
    if 'type' in data:
        edge_types[data['type']] += 1

for etype, count in sorted(edge_types.items()):
    print(f"{etype}: {count}")

# Find a specific function node to examine
print(f"\n=== Looking for Function Nodes with Code ===")
function_nodes = []
for node_id in G.nodes():
    node_data = G.nodes[node_id]
    if node_data.get('type') == 'function' and 'code_content' in node_data:
        function_nodes.append(node_id)
        if len(function_nodes) >= 3:
            break

if function_nodes:
    print(f"Found {len(function_nodes)} function nodes with code content")
    # Show one function example
    example_func = function_nodes[0]
    func_data = G.nodes[example_func]
    print(f"\nExample Function: {example_func}")
    print(f"  Start line: {func_data.get('start_line', 'N/A')}")
    print(f"  End line: {func_data.get('end_line', 'N/A')}")
    if 'code_content' in func_data and func_data['code_content']:
        code_lines = func_data['code_content'].split('\n')
        print(f"  Code length: {len(code_lines)} lines")
        print(f"  First 5 lines of code:")
        for i, line in enumerate(code_lines[:5], 1):
            print(f"    {i}: {line}")